#include "klib.h"
#include "vme.h"
#include "proc.h"

static TSS32 tss;

void init_gdt() {
  static SegDesc gdt[NR_SEG];
  gdt[SEG_KCODE] = SEG32(STA_X | STA_R,   0,     0xffffffff, DPL_KERN);
  gdt[SEG_KDATA] = SEG32(STA_W,           0,     0xffffffff, DPL_KERN);
  gdt[SEG_UCODE] = SEG32(STA_X | STA_R,   0,     0xffffffff, DPL_USER);
  gdt[SEG_UDATA] = SEG32(STA_W,           0,     0xffffffff, DPL_USER);
  gdt[SEG_TSS]   = SEG16(STS_T32A,     &tss,  sizeof(tss)-1, DPL_KERN);
  set_gdt(gdt, sizeof(gdt[0]) * NR_SEG);
  set_tr(KSEL(SEG_TSS));
}

void set_tss(uint32_t ss0, uint32_t esp0) {
  tss.ss0 = ss0;
  tss.esp0 = esp0;
}

typedef union free_page {
  union free_page *next;
  char buf[PGSIZE];
} page_t;

page_t *free_page_list;


static PD kpd;
static PT kpt[PHY_MEM / PT_SIZE] __attribute__((used));

void init_page() {
  extern char end;
  panic_on((size_t)(&end) >= KER_MEM - PGSIZE, "Kernel too big (MLE)");
  static_assert(sizeof(PTE) == 4, "PTE must be 4 bytes");
  static_assert(sizeof(PDE) == 4, "PDE must be 4 bytes");
  static_assert(sizeof(PT) == PGSIZE, "PT must be one page");
  static_assert(sizeof(PD) == PGSIZE, "PD must be one page");
  for (int i = 0; i < PHY_MEM / PT_SIZE; i++) {
    kpd.pde[i].val = MAKE_PDE(&kpt[i], 7);
    // kpd.pde[i].present = 1;
    // kpd.pde[i].write = 1;
    // kpd.pde[i].user = 1;
  }

  for (int i = 0; i < PHY_MEM / PT_SIZE; i++) {
    for (int j = 0; j < NR_PTE; j++) {
      kpt[i].pte[j].val = MAKE_PTE((i << DIR_SHIFT) | (j << TBL_SHIFT), 7);
      // kpt[i].pte[j].present = 1;
      // kpt[i].pte[j].write = 1;
      // kpt[i].pte[j].user = 1;
    }
  }

  // Lab1-4: init kpd and kpt, identity mapping of [0 (or 4096), PHY_MEM)
  // TODO();
  kpt[0].pte[0].val = 0;
  set_cr3(&kpd);
  set_cr0(get_cr0() | CR0_PG);
  // Lab1-4: init free memory at [KER_MEM, PHY_MEM), a heap for kernel
  free_page_list = (page_t *)(KER_MEM);

  page_t *tep = free_page_list;

  for (int i = KER_MEM / PGSIZE; i < PHY_MEM / PGSIZE - 1 ; i++) {
    tep++;
    // printf("tep: %p\n", tep);
    free_page_list->next = tep;
    free_page_list = tep;
  }
  
  //再把最后一个指向第一个
free_page_list = (page_t *)(KER_MEM);
}

void *kalloc() {
  // printf("free_page_list: %p\n", free_page_list);
  // Lab1-4: alloc a page from kernel heap, abort when heap empty
  if (free_page_list == NULL) {
    assert(0);
  }
  page_t *ret = free_page_list;
  free_page_list = free_page_list->next;
  return ret;

}

void kfree(void *ptr) {
  // Lab1-4: free a page to kernel heap
  // you can just do nothing :)
  //TODO();
  page_t *tep = (page_t *)ptr;
  tep->next = free_page_list;
  free_page_list = tep;
}

PD *vm_alloc() {
  // Lab1-4: alloc a new pgdir, map memory under PHY_MEM identityly
  PD *pgdir = kalloc();
  memcpy(pgdir, &kpd, sizeof(kpd));
  //其余的PDE清0
  //TODO();!!!!检查这里是不是有问题（感觉大概率没问题）
  for (int i = PHY_MEM / PT_SIZE; i < NR_PDE; i++) {
    pgdir->pde[i].val = 0;
  }
  
  return pgdir;
}

void vm_teardown(PD *pgdir) {
  // Lab1-4: free all pages mapping above PHY_MEM in pgdir, then free itself
  // you can just do nothing :)
  //TODO();
}

PD *vm_curr() {
  return (PD*)PAGE_DOWN(get_cr3());
}

PTE *vm_walkpte(PD *pgdir, size_t va, int prot) {
  //助教给的
  assert((prot & ~7) == 0);


  // Lab1-4: return the pointer of PTE which match va
  // if not exist (PDE of va is empty) and prot&1, alloc PT and fill the PDE
  // if not exist (PDE of va is empty) and !(prot&1), return NULL
  // remember to let pde's prot |= prot, but not pte

int pd_index = ADDR2DIR(va); // 计算“页目录号”
PDE *pde = &(pgdir->pde[pd_index]); // 找到对应的页目录项

if (pde->present == 0) { // 如果对应的页目录项不存在
  if (prot & PTE_P) { // 如果需要建立页表
    PT *pt = kalloc(); // 分配一个页表
    memset(pt, 0, sizeof(PT)); // 清零
    pde->val = MAKE_PDE(pt, prot); // 建立页目录项

    int pt_index = ADDR2TBL(va); // 计算“页表号”
    PTE *pte = &(pt->pte[pt_index]); // 找到对应的页表项

    
    return pte; // 返回页表项指针
  } else { // 如果不需要建立页表
    return NULL; // 返回空指针
  }
} else { // 如果对应的页目录项存在
  PT *pt = PDE2PT(*pde); // 根据PDE找页表的地址
int pt_index = ADDR2TBL(va); // 计算“页表号”
PTE *pte = &(pt->pte[pt_index]); // 找到对应的页表项

//辖写的，很可能错误
// pte->present = 1; // 设置为存在


return pte; // 返回页表项指针

}




//AI生成的不一定有用的代码
  // size_t dir = DIR(va);
  // size_t tbl = TBL(va);
  // if (pgdir->pde[dir].present == 0) {
  //   if (prot & PTE_P) {
  //     PT *pt = kalloc();
  //     memset(pt, 0, sizeof(PT));
  //     pgdir->pde[dir].val = MAKE_PDE(pt, prot);
  //   } else {
  //     return NULL;
  //   }
  // }
  // PTE *pte = &((PTE *)(pgdir->pde[dir].val & ~0xfff))[tbl];
  // return pte;



}

void *vm_walk(PD *pgdir, size_t va, int prot) {
  // Lab1-4: translate va to pa
  // if prot&1 and prot voilation ((pte->val & prot & 7) != prot), call vm_pgfault
  // if va is not mapped and !(prot&1), return NULL
  PTE *pte = vm_walkpte(pgdir, va, prot);
  if (pte == NULL) {
    return NULL;
  }


  /*权限相关的，AI生成的，不一定有用
  if ((pte->val & prot & 7) != prot) {
    vm_pgfault(va, 0);
  }*/
void *page = PTE2PG(*pte); // 根据PTE找物理页的地址
void *pa = (void*)((uint32_t)page | ADDR2OFF(va)); // 补上页内偏移量

return pa; // 返回物理地址
}

void vm_map(PD *pgdir, size_t va, size_t len, int prot) {
  // Lab1-4: map [PAGE_DOWN(va), PAGE_UP(va+len)) at pgdir, with prot
  // if have already mapped pages, just let pte->prot |= prot
  assert(prot & PTE_P);
  assert((prot & ~7) == 0);
  size_t start = PAGE_DOWN(va);
  size_t end = PAGE_UP(va + len);


  // 测试点
  // printf("start: %p, end: %p\n", start, end);
  // printf("%d", PHY_MEM);
  
  
  assert(start >= PHY_MEM);
  assert(end >= start);

  for (size_t cur = start; cur < end; cur += PGSIZE) {
    PTE *pte = vm_walkpte(pgdir, cur, prot);
    assert(pte != NULL);

    pte->val = MAKE_PTE(kalloc(), prot);
    
    
    
    
    //可以不实现的部分
    // if ((pte->val & prot & 7) != prot) {
    //   vm_pgfault(cur, 0);
    // }
    // pte->val |= prot;
  }

  
}

void vm_unmap(PD *pgdir, size_t va, size_t len) {
  // Lab1-4: unmap and free [va, va+len) at pgdir
  // you can just do nothing :)
  //assert(ADDR2OFF(va) == 0);
  //assert(ADDR2OFF(len) == 0);
  //TODO();
}

void vm_copycurr(PD *pgdir) {
  // Lab2-2: copy memory mapped in curr pd to pgdir
  size_t start = PHY_MEM;
  size_t end = USR_MEM;
  
  PD *cur_pg_dir = vm_curr();
  
  for (size_t va = start; va < end; va += PGSIZE) {
    PTE *pte = vm_walkpte(cur_pg_dir, va, 7);
    if (pte != NULL && pte->present) {
      
      vm_map(pgdir, va, PGSIZE, 7
      // pte->val & 0xfff
      );

      void *pa = vm_walk(pgdir, va, 7);
      
      void *cur_pa = vm_walk(cur_pg_dir, va, 7);

      
            // memcpy(pa, (void*)va, PGSIZE);

      memcpy(pa, cur_pa, PGSIZE);
      // break;
    }

  }

  
}

void vm_pgfault(size_t va, int errcode) {
  printf("pagefault @ 0x%p, errcode = %d\n", va, errcode);
  panic("pgfault");
}
